import numpy as np
from qiskit import QuantumCircuit, QuantumRegister, transpile
from qiskit.circuit.library import RYGate
from qiskit_aer import AerSimulator
from typing import Optional, List, Tuple

# 提高递归深度，避免树递归过深时报错（大向量构树时层数较多）
import sys
sys.setrecursionlimit(1000)


# =========================
# 二叉树结构：用于分治地存储向量幅值平方与相位
# =========================
class BinaryTree:
    def __init__(self, value, level: int = 0, branch: int = 0):
        """
        初始化二叉树节点。

        参数:
            value : 节点存储的数值
                - 对内部节点: 为该子树（对应一段子向量）的“幅值平方之和” s = Σ|x_i|^2
                - 对叶子节点: 存 [s, phase]，其中 s = |x|^2，phase = angle(x) 为该元素的相位
            level : 节点层级（根为 0，越往下越大）
            branch: 分支索引（可视为从根到该节点的二进制路径编码）
        """
        self.value = value
        self.child_left: Optional["BinaryTree"] = None
        self.child_right: Optional["BinaryTree"] = None
        self.level = level
        self.branch = branch

    # -------------------------
    # 插入左/右子节点（若已有，则封装一层，把旧子树下移）
    # -------------------------
    def insert_left(self, value, branch: int):
        """插入左孩子；若已存在则下移一层以保持插入为最上层。"""
        if self.child_left is None:
            self.child_left = BinaryTree(value, self.level + 1, branch)
        else:
            new_node = BinaryTree(value, self.level + 1, branch)
            # 下移旧的左子树（层级+1、分支索引重标）
            self.child_left.set_level(self.child_left.level + 1)
            self.child_left.set_branch(self.child_left.branch * 2)
            new_node.child_left = self.child_left
            self.child_left = new_node

    def insert_right(self, value, branch: int):
        """插入右孩子；若已存在则下移一层以保持插入为最上层。"""
        if self.child_right is None:
            self.child_right = BinaryTree(value, self.level + 1, branch)
        else:
            new_node = BinaryTree(value, self.level + 1, branch)
            # 下移旧的右子树（层级+1、分支索引重标）
            self.child_right.set_level(self.child_right.level + 1)
            self.child_right.set_branch(self.child_right.branch * 2 + 1)
            new_node.child_right = self.child_right
            self.child_right = new_node

    # -------------------------
    # 基本 getter / setter
    # -------------------------
    def get_value(self):
        """返回节点存储的值（内部节点为 s，总幅值平方；叶子为 [s, phase]）。"""
        return self.value

    def get_left(self) -> Optional["BinaryTree"]:
        return self.child_left

    def get_right(self) -> Optional["BinaryTree"]:
        return self.child_right

    def get_level(self) -> int:
        return self.level

    def get_branch(self) -> int:
        """返回分支索引（用于构造多控门时的控制比特选择）。"""
        return self.branch

    def set_level(self, level: int):
        self.level = level

    def set_branch(self, branch: int):
        self.branch = branch


def show_tree(T: Optional[BinaryTree]):
    """
    递归展示二叉树结构，便于调试。
    返回格式:
        ([value, level, branch], left_subtree, right_subtree)
    """
    if T is not None:
        return ([T.get_value(), T.get_level(), T.get_branch()],
                show_tree(T.get_left()),
                show_tree(T.get_right()))
    return None


# =========================
# 将向量转换为二叉树
# =========================
def create_tree_from_vector(x: List[complex]) -> Tuple[float, BinaryTree]:
    """
    根据输入向量构建一棵二叉树：
    - 对内部节点存储该节点覆盖的子段向量的幅值平方和 s
    - 对叶子节点存储 [s, phase]，其中 phase 是该元素的辐角

    这样做是为了分治地把幅值信息转为一系列旋转角（RY 的角度），
    同时在叶子层处理相位差（通过 P 相移）。
    """
    norm = np.linalg.norm(x)
    if norm == 0:
        raise ValueError("输入向量不能为零向量")
    u_x = [xi / norm for xi in x]  # 归一化

    # 长度为 1 的特例：根直接为叶
    if len(u_x) == 1:
        s = 1.0
        phase = np.angle(u_x[0])
        tree = BinaryTree([s, phase])
        return norm, tree

    # 一般情况：根为 1（总概率为 1），再递归构建左右子树
    tree = BinaryTree(1.0)
    u_x_left = u_x[: (len(x) + 1) // 2]
    u_x_right = u_x[(len(x) + 1) // 2 :]
    create_tree_recursively(u_x_left, tree, "left", tree.get_branch())
    create_tree_recursively(u_x_right, tree, "right", tree.get_branch())
    return norm, tree


def create_tree_recursively(x: List[complex], tree: BinaryTree, side: str, branch: int):
    """
    递归构建二叉树的辅助函数。

    逻辑:
    - 计算当前段的幅值平方和 s
    - 若为叶子（len(x)==1），则把 value 设为 [s, phase]
    - 将该 value 插入到父节点的左右孩子（由 side 指定）
    - 若非叶子，继续拆分成左右段，必要时补零（保证“近似完全二叉树”结构），并递归构建
    """
    # 当前子段的幅值平方和
    s = float(sum(abs(xi) ** 2 for xi in x))

    # 叶子节点：带相位
    if len(x) == 1:
        phase = float(np.angle(x[0]))
        s = [s, phase]  # 叶子存成 [s, phase]

    # 插入到父节点
    if side == "left":
        tree.insert_left(s, branch * 2)
    elif side == "right":
        tree.insert_right(s, branch * 2 + 1)
    else:
        raise ValueError("无效分支类型（side 需为 'left' 或 'right'）")

    # 继续递归构建孩子
    if len(x) > 1:
        mid = (len(x) + 1) // 2
        x_left = x[: mid]
        x_right = x[mid :]

        # 为了保持树更规整：如果一边是单元素而另一边不是，则在单元素一侧补一个 0
        if len(x_left) == 1 and len(x_right) > 1:
            x_left = [x_left[0], 0]
        elif len(x_right) == 1 and len(x_left) > 1:
            x_right = [x_right[0], 0]

        if side == "left":
            create_tree_recursively(x_left, tree.get_left(), "left", tree.get_left().get_branch())
            create_tree_recursively(x_right, tree.get_left(), "right", tree.get_left().get_branch())
        else:  # side == "right"
            create_tree_recursively(x_left, tree.get_right(), "left", tree.get_right().get_branch())
            create_tree_recursively(x_right, tree.get_right(), "right", tree.get_right().get_branch())


# =========================
# 根据二叉树把向量“加载”到量子电路
# =========================
def load_vector_from_qRAM(x: List[complex],
                          qc: Optional[QuantumCircuit] = None,
                          pos_init: int = 0) -> QuantumCircuit:
    """
    根据构建好的二叉树（幅值/相位的层次分解），在量子电路中实现向量加载。

    思想:
    - 根层通过一个 RY 在首个比特上设置“左右段”的幅值比例；
    - 之后逐层通过多控制 RY，把条件旋转分发到下层（控制比特由二进制路径决定）；
    - 叶子层若左右子元素存在非 {0, π} 的相位差，则用相位门 P(Δφ) 予以补偿。

    参数:
        x        : 输入复向量（会在内部归一化），长度不必是 2 的幂（内部会做“补零”处理）
        qc       : 可选的已有量子电路；若不传则内部创建
        pos_init : 在已有电路中，向量加载所使用的起始量子比特索引

    返回:
        构建好加载门序列的量子电路
    """
    norm, B_x = create_tree_from_vector(x)
    n = int(np.ceil(np.log2(len(x))))
    if n == 0:
        n = 1  # 长度为 1 的特例

    # 若未提供电路，则新建寄存器并创建电路（n 个量子比特）
    if qc is None:
        qr = QuantumRegister(n, "tree")
        qc = QuantumCircuit(qr)

    # 为了统一后续处理，先把 (pos_init ... pos_init+n-1) 的顺序翻转
    # 这是因为从高位到低位/从根到叶的控制序需要对齐
    for i in range(n // 2):
        qc.swap(pos_init + i, pos_init + n - 1 - i)

    # 取根的左右孩子（第一层）
    B_l = B_x.get_left()
    B_r = B_x.get_right()

    # 根层在第 pos_init 个量子比特上施加一个 RY:
    #   theta_0 = 2*arccos( sqrt( 概率落在“左子树”的比例 ) )
    # 若左孩子是叶子，则其 value 是 [s, phase]；否则为 s
    B_l_is_leaf = B_l.get_left() is None
    node_value_left = B_l.get_value()[0] if B_l_is_leaf else B_l.get_value()
    theta_0 = 2 * np.arccos(np.sqrt(node_value_left))
    qc.ry(theta_0, pos_init)

    # 处理根层叶子对的相位差：只有在左右两叶的相位差不等于 {0, π} 时，才需要用 P 门补偿
    is_complex_pair = False
    if B_l_is_leaf:
        phase_l, phase_r = B_l.get_value()[1], B_r.get_value()[1]
        # 四种“仅 ± 号差异”的特殊配置：不需要显式 P，只需把 RY 角度符号/偏置处理好即可
        if not ((phase_l, phase_r) in [(0, 0), (0, np.pi), (np.pi, 0), (np.pi, np.pi)]):
            is_complex_pair = True
        elif (phase_l, phase_r) == (0, np.pi):
            theta_0 = -theta_0
        elif (phase_l, phase_r) == (np.pi, 0):
            theta_0 = np.pi - theta_0
        elif (phase_l, phase_r) == (np.pi, np.pi):
            theta_0 = np.pi + theta_0

    if is_complex_pair:
        # 给当前位添加相位差（右 - 左），把相对相位“挂”到该位上
        qc.p(phase_r - phase_l, pos_init)

    # 若还有更深层（n>1），继续按二叉树递归构建多控制 RY 以细化到更小的子段
    if n > 1:
        processNode(qc, pos_init, n, B_l)
        processNode(qc, pos_init, n, B_r)

    qc.barrier()
    return qc


def processNode(q_circuit: QuantumCircuit, first_qubit: int, sizeRegisterTree: int, u_x: BinaryTree):
    """
    递归处理一个内部节点 u_x，向下对其左右孩子应用“多控制 RY”。
    - 控制比特由分支索引的二进制形式决定（哪条路径成立就触发旋转）
    - 旋转角由左右子段的概率占比决定：theta = 2*arccos( sqrt(s_left / s_parent) )
    - 若左右为叶子且相位差仅为 {0, π} 的组合，可折算进角度符号/偏置；否则需额外相位门（在父层已处置）

    参数:
        q_circuit       : 量子电路
        first_qubit     : 当前这组加载用的起始量子比特索引
        sizeRegisterTree: 本向量加载使用的量子寄存器大小（=n）
        u_x             : 当前内部节点（非叶）
    """
    u_l = u_x.get_left()
    u_r = u_x.get_right()
    if u_x.get_value() == 0 or u_l is None or u_r is None:
        # 若父段概率为 0 或孩子缺失，则无需处理
        return

    # 计算该层 RY 的旋转角：由“左子段占父段的概率比例”确定
    u_l_is_leaf = u_l.get_left() is None
    node_value_left = u_l.get_value()[0] if u_l_is_leaf else u_l.get_value()
    theta = 2 * np.arccos(np.sqrt(node_value_left / u_x.get_value()))

    # 如果左右为叶子，检查是否需要相位补偿（仅 {0, π} 组合可以通过符号调整消化）
    if u_l_is_leaf:
        phase_l, phase_r = u_l.get_value()[1], u_r.get_value()[1]
        if (phase_l, phase_r) == (0, np.pi):
            theta = -theta
        elif (phase_l, phase_r) == (np.pi, 0):
            theta = np.pi - theta
        elif (phase_l, phase_r) == (np.pi, np.pi):
            theta = np.pi + theta
        # 若出现一般相位差（非 {0,π}），在父层/根层会用 P 门处理，这里无需额外动作

    # 依据分支索引，决定哪些控制比特应为 1（否则先打 X 把 0 变 1，以便统一用“全为1”触发）
    list_binary = getBinary(u_x.get_branch(), u_l.get_level() - 1)[0]  # 只取整数部分的二进制编码
    for i, b in enumerate(list_binary):
        if b == '0':
            q_circuit.x(first_qubit + i)

    # 构建多控制 RY（控制数 = u_l 的层级数 - 1；目标为该层位）
    # 约定: 当前层的目标位是 first_qubit + (u_l.get_level() - 1)
    # 但这里和原逻辑保持一致：使用 [first_qubit + i for i in range(u_l.get_level())]
    # 其中最后一个为目标，其余为控制（Qiskit 的 .control(k) 约定：前 k 位是控制，最后接目标）
    MC_Ry = RYGate(theta).control(u_l.get_level() - 1)
    q_circuit.append(MC_Ry, [first_qubit + i for i in range(u_l.get_level())])

    # 把先前打的 X 还原（恢复计算基态）
    for i, b in enumerate(list_binary):
        if b == '0':
            q_circuit.x(first_qubit + i)

    # 若左孩子仍是内部节点，则继续递归细化到更小的段
    if not u_l_is_leaf:
        q_circuit.barrier()
        processNode(q_circuit, first_qubit, sizeRegisterTree, u_l)
        processNode(q_circuit, first_qubit, sizeRegisterTree, u_r)


# =========================
# 工具函数
# =========================
def getBinary(number: int, integerPrecision: int = 1) -> Tuple[str, str]:
    """
    把非负整数 number 转为二进制字符串（整数部分），附带伪“小数”二进制（这里实际不会用到）。
    返回 (binary_int, binary_dec)

    参数:
        number           : 非负整数
        integerPrecision : 最少位宽（不足则左侧补 0）
    """
    int_part = int(number)
    binary_int = format(int_part, '#0' + str(2 + integerPrecision) + 'b')[2:]  # 去掉 '0b' 前缀

    # 对于我们的使用场景，branch 是整数，小数部分恒为空，这段仅为完整性保留
    dec_part = number - int_part
    binary_dec = ''
    while dec_part > 0:
        dec_part *= 2
        if dec_part >= 1:
            binary_dec += '1'
            dec_part -= 1
        else:
            binary_dec += '0'
    return binary_int, binary_dec


def normalize_vector(v: List[complex]) -> np.ndarray:
    """
    把 Python 列表转为复数 numpy 向量，并做 L2 归一化。
    """
    arr = np.asarray(v, dtype=np.complex128)
    norm = np.linalg.norm(arr)
    if norm == 0:
        raise ValueError("向量不能全为零")
    return arr / norm


def simulate_and_verify(qc: QuantumCircuit, target_v: List[complex]) -> np.ndarray:
    """
    使用 Aer 的 statevector 模拟器执行电路，读取最终态向量并与目标向量计算保真度。

    流程:
        1) 在电路末尾保存 statevector
        2) 使用 AerSimulator(method='statevector') 运行
        3) 取出状态向量，与目标态（归一化后）做内积求 |<ψ|φ>|^2

    注意:
        - 此处的目标态向量应与电路比特顺序一致（本实现对位序做了 swap 以对齐）
    """
    qc.save_statevector()
    simulator = AerSimulator(method='statevector')
    tqc = transpile(qc, simulator)
    job = simulator.run(tqc)
    result = job.result()
    statevector = result.get_statevector()

    target_state = normalize_vector(target_v)
    fidelity = np.abs(np.vdot(target_state, statevector)) ** 2  # vdot 含共轭
    print(f"模拟状态向量: {np.real(statevector)}")
    print(f"目标状态: {target_state}")
    print(f"保真度: {fidelity:.6f}")
    return statevector


# =========================
# 最小可运行示例
# =========================
if __name__ == "__main__":
    # 示例：加载长度为 2 的复向量 [1, 1j]
    b = [1, 3, 0, 4]

    # 构建加载电路
    qc = load_vector_from_qRAM(x=b)

    # 模拟并核对保真度（应接近 1）
    simulate_and_verify(qc, b)
